/**
 CollinsKt.h

 \file
 Declaration of class CollinsKt.

 \author Thomas Burton 
 \date 2/28/12
 \copyright 2012 BNL. All rights reserved.
*/

#ifndef _CollinsKt_H_
#define _CollinsKt_H_
#include <TAxis.h>
#include <map>
#include <sstream>
#include <string>
#include <utility>
#include <vector>
#include <list>

#ifndef __CINT__
#include <boost/shared_ptr.hpp>
#endif

#include <Rtypes.h>

#include "tppmc/Bins.h"
#include "tppmc/DssFragmentation.h"
#include "tppmc/Enumerations.h"
#include "tppmc/FragmentationFunction.h"
#include "tppmc/Particle.h"
#include "tppmc/TF2J.h"

namespace tppmc {

   namespace collins {

      /**
       \todo Move implementation to .cxx file
      */
      class FragBin {
         
      public:
         
         int favoured; ///< favoured, unfavoured, or the average of those
         double depolMin;
         double depolMax;
         
         bool operator<(const FragBin& other) const {
            if(favoured < other.favoured) return true;
            if(favoured > other.favoured) return false;
            // favoured is equal, compare on depol range
            if(depolMin < other.depolMin) return true;
            if(depolMax > other.depolMax) return false;
            return depolMax < other.depolMax;
         }
         void Speak() const {
            std::cout << favoured << " " << depolMin << " " << depolMax << std::endl;
         }
         bool Match(const Hadron& hadron) const {
            int fav = 0; // By default assume unfavoured
            int htype = hadron.type;
            int ptype = hadron.mParton.type;
            switch(htype) {
               case PIPLUS:
                  if(2 == ptype or -1 == ptype) fav = 1;
                  break;
               case PIMINUS:
                  if(1 == ptype or -2 == ptype) fav = 1;
                  break;
               case PIZERO: // pi0 from u(bar), d(bar) is average --> fav = 2
                  if(2 == ptype or -1 == ptype or
                     -2 == ptype or 1 == ptype) fav = 2;
                  break;
            } // switch
            if(fav not_eq favoured) return false;
            if(hadron.mParton.cos_spin_depol < depolMin or
               hadron.mParton.cos_spin_depol >= depolMax) return false;
            return true;
         }               
      };

      class FunctorNew; // Forward declaration so FunctionTable can create
                        // a pointer list
      /**
       The collection of ROOT functions used by PositivityKt.
       The function definitions are all the same, but have different parameters.
       A different copy of the function is maintained for each permutation of
       parameters to avoid recomputation of the function integral each time the
       parameters change, which gets computationally expensive.
      */
      class FunctionTable {
      public:

         FunctionTable();

#ifndef __CINT__
         typedef boost::shared_ptr<tppmc::FragmentationFunction> FragFcnPtr;
         FunctionTable(FragFcnPtr);
#endif

         /** Destructor */
         virtual ~FunctionTable();

         /**
          Returns the ROOT wrapper around the fragmentation function
          for computing fragmentation for a hadron at this scale.
          The hadron's mParton attribute must be set and define type and
          cos_spin_depol.
          Returns NULL if there is no function for this parton/hadron and Q bin.
          The pointer is non-const because methods like GetRandom() and
          Integral() are non-const, but don't modify the function e.g. by
          changing parameters.
         */
         TF2* Find(const Hadron&, double /*Q2*/);

         /**
          Populate the table with functions for each bin in depol and Q.
          Functions are evaluated at the lower edges of the bins.
         */
         int Initialise(const Bins* depol, const Bins* Q);

         void Speak();
         
      protected:
#ifndef __CINT__
         TAxis* mQ;
         typedef std::vector<TF2*> TF2Vec;
         typedef std::map<FragBin, TF2Vec> Table;
         std::map<FragBin, TF2Vec> mFunctions;
         TF2* BuildFunction(FragBin, double Q);
         FragFcnPtr mUnpol;
         std::list<FunctorNew*> mFunctors;
#endif
         ClassDef(tppmc::collins::FunctionTable, 1)
      };


      /**
       kT-dependent positivity bound based on a kT-integrated unpolarised
       fragmentation function.
       Only pion production is currently supported.
       The positivity bound is defined via
       |collins(z, pt)| = 2 * unpolarised(z, pt)
       where collins(z, pt) is the Anselmino et al. definition of the
       Collins function, such that
       fftotal(z, pt) = unpolarised(z, pt) + 0.5s.(q x pt)collins(z, pt)
       where s is the parton spin and q the parton 3-momentum.
       As the pt dependence is not well known, it is factorised according to
       unpolarised(z, pt) = unpolarised(z) * exp(-pt2/<pt2>) / pi / <pt2>
       pt2 is the squared hadron pt wrt parton 3-momentum direction and
       <pt2> is the mean pt2. This is settable as a parameter (default is
       0.2 GeV^2). The unpolarised fragmentation can be set via the constructor.
       The default is leading order DSS.
       \todo Move headers for code in collins namespace to tppmc/collins/
       subdirectory, but check impact on LinkDef/CINT dictionary first.
       */
      class PositivityKt : public FragmentationFunction {

      public:

         /**
          Initialises with default distribution (LO DSS) and mean pt2 = 0.2
          */
         PositivityKt();

#ifndef __CINT__
         PositivityKt(boost::shared_ptr<FragmentationFunction>,
                      double meanPt2);
#endif

         virtual ~PositivityKt();


         virtual bool Initialise();


         /**
          Fragment the input parton to a single hadron at the requested scale.
          If the parton species is not supported, fragment via the normal
          DSS fragmentation function, with a random kT according to a
          Gaussian distribution.
          */
         virtual std::vector<tppmc::Hadron> Fragment(const Parton&,
                                                     double scale);

         /**
          Evaluate the fragmentation function for the input parton to
          fragment to the input hadron at the requested scale.
          */
         virtual double Evaluate(const Parton&, const Hadron&, double scale);

         /**
          Returns true if the fragmentation of the input parton type
          is supported by this function.
          */
         virtual bool Supports(const Parton&) const;
         
         /**
          Returns true if the fragmentation of the input parton type
          is supported by this function.
          */
         virtual bool Supports(const Hadron&) const;
         
         /**
          Returns the name of this fragmentation function
          */
         virtual std::string Name() const;

         virtual void SetHadron(int = INVALID_HADRON);

         virtual void SetZRange(double lower, double upper);

         /**
          Lists hadron species supported by the fragmentation function.
          */
         virtual const std::vector<int>& Hadrons() const;

         /**
          Lists parton species supported by the fragmentation function.
          */
         virtual const std::vector<int>& Partons() const;

         virtual int ThrowHadronType(const Parton& parton, double Q2);

         /** Returns the minimum z for which the distribution is valid */
         virtual double ZMin() const;
         
         /** Returns the maximum z for which the distribution is valid */
         virtual double ZMax() const;
         
         /**
          The lowest fragmentation scale allowed by the fragmentation function.
          */
         virtual double MinimumScale() const;
         
         /**
          The highest fragmentation scale allowed by the fragmentation function.
          */
         virtual double MaximumScale() const;

         virtual void Qa(std::string filename);

      protected:

         void ComputeFragmentationProbabilities();
         double mMeanPt2;
         
         std::vector<int> mPartons; ///< Supported parton PDG codes
         std::vector<int> mHadrons; ///< Supported hadron PDG codes
#ifndef __CINT__
         typedef boost::shared_ptr<FragmentationFunction> FragFcnPtr;
         FragFcnPtr mUnpol; //!< Unpolarised
                                                          ///< frag function

         typedef boost::shared_ptr<FragmentationTable> FragTablePtr;
         std::map<int, FragTablePtr> mProbabilities; //!< Probability for each
                                                     ///< parton to fragment
                                                     ///< each hadron at each Q
         /**\todo change to smart pointer*/
         Bins* mQBins; //!< Bins in Q: NOT Q-squared!
         
         FunctionTable mFunctions;//!
#endif
         ClassDef(tppmc::collins::PositivityKt, 1)
      };

      /**
       \todo Move to its own file
      */
      class KtDependent {
         
      public:
         
         /**
          Information describing parameters 
          */
         class FragInfo {
            
         public:
            
            int favoured;
            double depolMin;
            double depolMax;
            
            bool operator<(const FragInfo& other) const {
               if(favoured < other.favoured) return true;
               if(favoured > other.favoured) return false;
               // favoured is equal, compare on depol range
               if(depolMin < other.depolMin) return true;
               if(depolMax > other.depolMax) return false;
               return depolMax < other.depolMax;
            }
            void Speak() const {
               std::cout << favoured << " " << depolMin << " " << depolMax << std::endl;
            }
            bool Match(const Hadron& hadron) const {
               int fav = -1;
               int htype = hadron.type;
               int ptype = hadron.mParton.type;
               switch(htype) {
                  case PIPLUS:
                     if(2 == ptype or -1 == ptype) fav = 1;
                     else fav = 0;
                     break;
                  case PIMINUS:
                     if(1 == ptype or -2 == ptype) fav = 1;
                     else fav = 0;
                     break;
                  case PIZERO:
                     if(2 == ptype or -1 == ptype or
                        -2 == ptype or 1 == ptype) fav = 1;
                     else fav = 0;
                     break;
               } // switch
               std::cout << "Matching: parton " << ptype << " hadron " << htype << " fav = " << fav << std::endl;
               if(fav not_eq favoured) return false;
               
               if(hadron.mParton.cos_spin_depol < depolMin or
                  hadron.mParton.cos_spin_depol >= depolMax) return false;
               return true;
            }               
            
            ClassDef(tppmc::collins::KtDependent::FragInfo, 1)
         };
         
         
         typedef std::map<FragInfo, std::string> Table;
         Table mTable;
         
         KtDependent() {
            
            double depol[11] = {
               0., 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.
            };
            std::stringstream ss;
            for(int i(0); i < 2; ++i) {
               for(int j(0); j < 10; ++j) {
                  FragInfo f;
                  f.favoured = i;
                  f.depolMin = depol[j];
                  f.depolMax = depol[j + 1];
                  ss.str(""); ss.clear();
                  ss << "Favoured " << f.favoured << " dmin = " <<
                  f.depolMin << " dmax = " << f.depolMax;
                  mTable.insert(std::make_pair(f, ss.str()));
               } // for
            } // for
         }
         
         std::string Find(const Hadron& hadron) const;
         
         ClassDef(tppmc::collins::KtDependent, 1)
      };
      
   } // namespace collins
   
} // namespace tppmc

#endif
